#################################
# Balance test table function
#
# Replication Material for:
# Continuity or Change? (In)direct Rule in British and French Colonial Africa
# 
# Carl Mueller-Crepon, 2020
# International Organization
#
# File Description:
# Function to produce a balance test table from a set of models.
#
# Called from the main analysis files
#
################################

# Round number to digits
specify_decimal <- function(x, digits){format(round(x, digits), nsmall = digits)}

# Standard error to string
get_se_string <- function(point, se, p, digits = 2){
  se.str <- paste0("(",specify_decimal(se, digits = digits),")")
  if(p < 0.01){
    point.str <- paste0(specify_decimal(point, digits = digits),"^{***}")
  } else if(p<0.05){
    point.str <- paste0(specify_decimal(point, digits = digits),"^{**}")
  } else if(p<0.1){
    point.str <- paste0(specify_decimal(point, digits = digits),"^{*}")
  } else{
    point.str <- paste0(specify_decimal(point, digits = digits))
  }
  return(list(point.str, se.str))
}

# Balance table as LaTeX Table
balance_table_tex <- function(model.ls = NULL, coef.mat = NULL, se.mat = NULL, p.mat = NULL, 
                              treat.var = NULL, treat.var.lab = treat.var,
                              balance.var.label = NULL,
                              digits = 3, stars = T, std.coef = T,
                              label = "", col.sep.width = "5pt", title = "Balance test", 
                              col.header = "", colhead.sep = NULL, notes = "", add.rows = NULL, indep.head.lab = "Indep. variable",
                              cmidrule = F, cmidrule.trim = 0){
  if(length(treat.var) == 1){
    treat.var = rep(treat.var, length(model.ls))
    treat.var.lab = rep(treat.var.lab, length(model.ls))
  }
  if(is.null(colhead.sep)){
    colhead.sep <- rep(1, length(col.header))
  }
  # Get coefficients and SEs
  if(is.null(coef.mat)){
    coef.mat <- do.call(cbind, lapply(c(1:length(model.ls)), function(i){
      unlist(lapply(model.ls[[i]], function(m){
        if(std.coef){
          m$coefficients[treat.var[i],1] / sd(m$response)
        } else {
          m$coefficients[treat.var[i],1]
        }
      }))
    }))
  }
  if(is.null(se.mat)){
    se.mat <- do.call(cbind, lapply(c(1:length(model.ls)), function(i){
      unlist(lapply(model.ls[[i]], function(m){
        if(std.coef){
          m$clustervcv[treat.var[i],treat.var[i]]^.5 / sd(m$response)
        } else {
          m$clustervcv[treat.var[i],treat.var[i]]^.5
        }
      }))
    }))
  }
  if(is.null(p.mat)){
    p.mat <- do.call(cbind, lapply(c(1:length(model.ls)), function(i){
      unlist(lapply(model.ls[[i]], function(m){
        m$cpval[treat.var[i]]
      }))
    }))
  }
 
  
  # ... var labels
  if(is.null(balance.var.label)){
    balance.var.label <- unlist(lapply(model.ls[[1]], function(m){
      m$lhs
    }))
  }
  
  # Print latex
  
  # ... column header
  col.head <- paste(paste0("\\multicolumn{",colhead.sep,"}{c}{",col.header, "}"), collapse = " & ")
  if(cmidrule){
    cmidrule <- paste(paste0("\\cmidrule(lr{",cmidrule.trim,"pt}){", c(2, 2 + cumsum(colhead.sep[-length(colhead.sep)])), "-",
                             c(1+ cumsum(colhead.sep)), "}"), collapse = " ")
  } else {
    cmidrule <- " \\hline "
  }
  

  # ... header
  latex.head <- paste0("\\begin{table}[!htbp] \\centering \\caption{", title, 
                       "} \\label{", label, "}  \\begin{tabular}{@{\\extracolsep{",col.sep.width,"}}l",
                       paste(rep(paste0("D{.}{.}{-", digits,"} "), length(model.ls)), collapse = "")," } 
                       \\\\[-1.8ex]\\hline ",
                       "\\hline \\\\[-1.8ex] ",
                       "\\\\[-1.8ex] \\multicolumn{1}{c}{} & ",
                       col.head," \\\\ ",cmidrule, 
                       " \\\\[-1.8ex]  \\multicolumn{1}{c}{",indep.head.lab,"} & ",
                       paste(paste0("\\multicolumn{1}{c}{",treat.var.lab, "}"), collapse = " & ")," \\\\ ",
                       "\\multicolumn{1}{l}{Dep. variable} &  &  &  \\\\ ",
                       "\\hline \\\\[-1.8ex] ")
  
  # ... content
  latex.body <- paste(unlist(lapply(c(1:nrow(coef.mat)), function(r){
    # .. get strings
    s.ls <- lapply(1:ncol(coef.mat), function(c){
      get_se_string(point = coef.mat[r,c], se = se.mat[r,c], 
                    p = p.mat[r,c], digits = digits)
    })
    # .. coef
    c.t <- paste(c(balance.var.label[r], 
                 unlist(lapply(s.ls, function(x){x[[1]]}))),
                 collapse = " & ")
    # .. SE
    s.t <- paste(c("", 
                   unlist(lapply(s.ls, function(x){x[[2]]}))),
                 collapse = " & ")
    
    # ... combine
    paste(c.t, " \\\\ \n", s.t)
  })), 
  collapse = paste("\\\\ \\\\[-1.8ex] \n ")) #, 
                   # paste(rep(" &", length(model.ls)), collapse = ""), 
                   # " \\\\ \n "))
  
  # ... latex end
  latex.foot <- paste0("\\hline \\hline  \\\\[-1.8ex]  ",
                       "\\multicolumn{",ncol(coef.mat)+1,"}{l}{", 
                       notes, "} \\end{tabular} \\end{table} ")
  
  # Combine
  latex.final <- paste(latex.head, " \\\\ \n ",latex.body)
  if(!is.null(add.rows)){
    latex.final <- paste(latex.final, " \\\\ \\hline \\\\[-1.8ex] \n  ",
                         paste(unlist(lapply(add.rows, function(r){
                           paste(r, collapse = " & ")
                         })), collapse = " \\\\ \n "))
  }
  latex.final <- paste(latex.final, " \\\\ \n ",latex.foot)
  
  # Return
  return(latex.final)
}


